#encoding=utf-8
from openai import OpenAI
from api.global_config import GPT3_CONFIG
from util.post_processing import extract_sentences, get_stripped_sentences, find_word_in_string
import string

class GPT3:
    def __init__(self, config: GPT3_CONFIG):
        self.config = config

    def make_requests(self, prompt):
        # print(f"Calling GPT-3.5 API with config: {self.config.__dict__}")
        client = OpenAI(
            api_key=self.config.api_key,
        )
        response = client.completions.create(
            model=self.config.engine,
            prompt=prompt,
            max_tokens=self.config.max_tokens,
            temperature=self.config.temperature,
            top_p=self.config.top_p
        )
        return response

    def post_processing(self, response):
        if response is None or response.choices[0].finish_reason=="length":
            return []
        raw_instructions = get_stripped_sentences(extract_sentences(response.choices[0].text))
        instructions = []
        for inst in raw_instructions:
            if inst == "":
                continue
            # filter out too short or too long instructions
            if len(inst) <= 3 or len(inst) > 150:
                continue
            # filter based on keywords that are not suitable for language models.
            if any(find_word_in_string(word, inst) for word in
                   ["image", "images", "graph", "graphs", "picture", "pictures", "file", "files", "map", "maps", "draw",
                    "plot", "go to"]):
                continue
            # We found that the model tends to add "write a program" to some existing instructions, which lead to a lot of such instructions.
            # And it's a bit confusing whether the model need to write a program or directly output the result.
            # Here we filter them out.
            # Note this is not a comprehensive filtering for all programming instructions.
            if inst.startswith("Write a program"):
                continue
            # filter those starting with punctuation
            if inst[0] in string.punctuation:
                continue
            # filter those starting with numbers
            if inst[0].isnumeric():
                continue
            instructions.append(inst)

        return instructions, response

    def get_raw_output(self, prompt):
        return self.make_requests(prompt).choices[0].text

    def get_instructions(self, prompt):
        return self.post_processing(self.make_requests(prompt))


if __name__ == "__main__":
    GPT3_API_KEY = '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'
    GPT3_ENGINE = 'gpt-3.5-turbo-instruct'

    prompt="Come up with a series of tasks:1.What's your name? 2.How's the weather? 3.How old are you? 4."
    prompt_chinese = 'Come up with a series of instructions with the same format, no more than 15 instructions:\n1. 如何使用信息技术为思政教育赋能？\n2. 请详细说明劳模精神、劳动精神、工匠精神的内涵和作用分别是什么，以及如何培育这些精神\n3. 红蓝融合中，理论和实践哪一个更重要？\n4. 见义勇为是否值得提倡？\n5. 信息时代，如何激发创新活力，产出创新成果？\n6. 请列举几个知行合一的名人事例\n7. 领导让我作一周工作的汇报总结，但我这周没有做实质性的工作，我该怎么办？\n8. 为什么要讲红蓝融合，其内涵是什么？\n9.'
    # 对于补全模型，很容易超字数，因此要限制最多15条指令，其中已经有了8条

    # 自己配置config
    config = GPT3_CONFIG(
        engine=GPT3_ENGINE,
        max_tokens=1024,
        temperature=0.7,
        top_p=0.5,
        frequency_penalty=0,
        presence_penalty=2,
        stop_sequences=["\n\n", "\n16", "16.", "16 ."],
        logprobs=1,
        n=1,
        best_of=1,
        api_key=GPT3_API_KEY
    )
    # 也可以使用默认config
    config_default = GPT3_CONFIG()
    # 使用配置对象创建GPT3类的实例
    gpt3_instance = GPT3(config_default)

    # 给出prompt，调用API
    result, _ = gpt3_instance.get_instructions(prompt_chinese)
    print(result)
    print(_)